# Copyright (c) Facebook, Inc. and its affiliates.
# -*- coding: utf-8 -*
import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import os.path as osp
import torch
import tkinter as tk
from tkinter import filedialog, ttk
import cv2
from datetime import datetime
from tkinter import ttk

# FrankMocap模块导入
from demo.demo_options import DemoOptions
from bodymocap.body_mocap_api import BodyMocap
from handmocap.hand_mocap_api import HandMocap
from bodymocap.body_bbox_detector import BodyPoseEstimator
from handmocap.hand_bbox_detector import HandBboxDetector
from renderer.visualizer import Visualizer
import mocap_utils.demo_utils as demo_utils

# 导入三种模式的运行函数
from demo.demo_handmocap import run_hand_mocap
from demo.demo_bodymocap import run_body_mocap
from demo.demo_frankmocap import run_frank_mocap

class UnifiedMocapApp:
    def __init__(self, root):
        self.root = root
        self.root.title("FrankMocap")
        self.root.geometry("800x650")  # 增加窗口尺寸
        self.root.configure(bg="#f0f0f0")
        
        # 设置支持中文的字体
        self.chinese_font = ("SimSun", 12)  # 微软雅黑，支持中文
        
        # 创建放大的样式
        self.create_styles()
        self.setup_ui()
        
    def create_styles(self):
        """创建自定义放大样式"""
        style = ttk.Style()
        style.configure("TLabel", font=self.chinese_font)
        style.configure("TRadiobutton", font=self.chinese_font)
        style.configure("TButton", font=self.chinese_font)
        style.configure("TEntry", font=self.chinese_font)
        style.configure("TCombobox", font=self.chinese_font)
        style.configure("TLabelframe", font=(self.chinese_font[0], 12, "bold"))
        style.configure("TLabelframe.Label", font=(self.chinese_font[0], 12, "bold"))
        style.configure("Run.TButton", font=(self.chinese_font[0], 14, "bold"))
    def setup_ui(self):
        # 标题标签
        title_label = tk.Label(
            self.root, 
            text="FrankMocap", 
            font=("Arial", 24, "bold"),  # 放大标题字体
            bg="#f0f0f0"
        )
        title_label.pack(pady=30)
        
        # 主框架
        main_frame = ttk.Frame(self.root)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=30, pady=15)  # 增加边距
        
        # 捕捉模式选择
        mode_frame = ttk.LabelFrame(main_frame, text="选择捕捉模式")
        mode_frame.pack(fill=tk.X, pady=15)
        
        self.mode_var = tk.StringVar(value="frank")
        modes = [
            ("完整捕捉 (身体+手部)", "frank"),
            ("仅身体捕捉", "body"),
            ("仅手部捕捉", "hand")
        ]
        
        for i, (text, value) in enumerate(modes):
            ttk.Radiobutton(
                mode_frame, 
                text=text,
                value=value,
                variable=self.mode_var
            ).pack(anchor=tk.W, padx=30, pady=8)  # 增加内边距
        
        # 视频选择框架
        video_frame = ttk.LabelFrame(main_frame, text="输入视频")
        video_frame.pack(fill=tk.X, pady=15)
        
        video_input_frame = ttk.Frame(video_frame)
        video_input_frame.pack(fill=tk.X, padx=10, pady=10)
        
        self.video_path_var = tk.StringVar()
        ttk.Entry(video_input_frame, textvariable=self.video_path_var, font=("Arial", 12)).pack(side=tk.LEFT, fill=tk.X, expand=True)
        ttk.Button(video_input_frame, text="浏览...", command=self.browse_video).pack(side=tk.RIGHT, padx=10)
        
        # 示例视频选择
        sample_frame = ttk.Frame(video_frame)
        sample_frame.pack(fill=tk.X, padx=10, pady=10)
        
        # 原代码
        # ttk.Label(sample_frame, text="或选择示例视频:").pack(side=tk.LEFT)

        # 修改后的代码
        label = ttk.Label(sample_frame, text="示例:", width=15)
        label.pack(side=tk.LEFT, padx=(0, 5))
        
        self.sample_var = tk.StringVar()
        samples = {
            "手部示例": "./sample_data/han_hand_short.single_totalbody.mp4",
            "身体示例": "./sample_data/han_short.mp4",
            "全身示例": "./sample_data/single_totalbody.mp4"
        }
        
        sample_combo = ttk.Combobox(sample_frame, textvariable=self.sample_var, values=list(samples.keys()), state="readonly", font=("Arial", 12))
        sample_combo.pack(side=tk.LEFT, padx=10)
        sample_combo.config(width=15)  # 设置下拉框宽度
        sample_combo.bind("<<ComboboxSelected>>", lambda e: self.select_sample(samples[self.sample_var.get()]))
        
        # 输出目录框架
        output_frame = ttk.LabelFrame(main_frame, text="输出设置")
        output_frame.pack(fill=tk.X, pady=15)
        
        output_dir_frame = ttk.Frame(output_frame)
        output_dir_frame.pack(fill=tk.X, padx=10, pady=10)
        
        ttk.Label(output_dir_frame, text="地址:").pack(side=tk.LEFT)
        self.output_path_var = tk.StringVar(value="./mocap_output")
        ttk.Entry(output_dir_frame, textvariable=self.output_path_var, font=("Arial", 12)).pack(side=tk.LEFT, fill=tk.X, expand=True, padx=10)
        ttk.Button(output_dir_frame, text="浏览...", command=self.browse_output_dir).pack(side=tk.RIGHT)
        
        # 渲染器选择
        renderer_frame = ttk.Frame(output_frame)
        renderer_frame.pack(fill=tk.X, padx=10, pady=10)
        
        ttk.Label(renderer_frame, text="渲染器:").pack(side=tk.LEFT)
        self.renderer_var = tk.StringVar(value="pytorch3d")
        renderer_combo = ttk.Combobox(
            renderer_frame, 
            textvariable=self.renderer_var,
            values=["pytorch3d", "opendr", "mesa"],
            state="readonly",
            font=("Arial", 12)
        )
        renderer_combo.pack(side=tk.LEFT, padx=10)
        renderer_combo.config(width=12)  # 设置下拉框宽度
        
        # 运行按钮
        run_button = ttk.Button(
            main_frame,
            text="Launch",
            command=self.run_mocap,
            style="Run.TButton"
        )
        run_button.pack(pady=30)
        
        # 状态标签
        self.status_var = tk.StringVar(value="准备就绪")
        status_label = ttk.Label(self.root, textvariable=self.status_var, font=("Arial", 12))
        status_label.pack(side=tk.BOTTOM, fill=tk.X, padx=15, pady=15)
    
    def browse_video(self):
        file_path = filedialog.askopenfilename(
            title="选择视频文件",
            filetypes=[
                ("视频文件", "*.mp4 *.avi *.mov *.mkv"),
                ("所有文件", "*.*")
            ]
        )
        if file_path:
            self.video_path_var.set(file_path)
            
    def select_sample(self, path):
        self.video_path_var.set(path)
    
    def browse_output_dir(self):
        dir_path = filedialog.askdirectory(title="选择输出目录")
        if dir_path:
            self.output_path_var.set(dir_path)
            
    def run_mocap(self):
        mode = self.mode_var.get()
        video_path = self.video_path_var.get()
        output_path = self.output_path_var.get()
        renderer_type = self.renderer_var.get()
        
        if not video_path:
            self.status_var.set("错误: 请选择视频文件")
            return
            
        if not os.path.exists(video_path):
            self.status_var.set(f"错误: 视频文件不存在: {video_path}")
            return
            
        # 创建输出目录
        if not os.path.exists(output_path):
            os.makedirs(output_path, exist_ok=True)
            
        # 更新状态
        self.status_var.set(f"正在处理: {os.path.basename(video_path)} 使用{mode}模式...")
        self.root.update()
        
        # 处理命令行参数
        sys.argv = [sys.argv[0]]
        sys.argv.extend(["--input_path", video_path])
        sys.argv.extend(["--out_dir", output_path])
        sys.argv.extend(["--renderer_type", renderer_type])
        
        # 运行相应的模式
        try:
            # 解析命令行参数
            args = DemoOptions().parse()
            args.use_smplx = True  # 启用SMPLX模型
            
            # 设置设备
            device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
            if not torch.cuda.is_available():
                self.status_var.set("警告: CUDA不可用，性能可能受限")
            
            # 创建可视化器
            if args.renderer_type in ['pytorch3d', 'opendr']:
                from renderer.screen_free_visualizer import Visualizer
            else:
                from renderer.visualizer import Visualizer
            visualizer = Visualizer(args.renderer_type)
            
            # 根据选择的模式运行不同的捕捉功能
            if mode == "hand":
                # 手部捕捉
                hand_bbox_detector = HandBboxDetector(args.view_type, device)
                hand_mocap = HandMocap(args.checkpoint_hand, args.smpl_dir, device=device)
                run_hand_mocap(args, hand_bbox_detector, hand_mocap, visualizer)
                
            elif mode == "body":
                # 身体捕捉
                body_bbox_detector = BodyPoseEstimator()
                checkpoint_path = args.checkpoint_body_smplx if args.use_smplx else args.checkpoint_body_smpl
                body_mocap = BodyMocap(checkpoint_path, args.smpl_dir, device, args.use_smplx)
                run_body_mocap(args, body_bbox_detector, body_mocap, visualizer)
                
            else:  # frank (默认)
                # 完整捕捉 (身体+手部)
                hand_bbox_detector = HandBboxDetector('third_view', device)
                body_mocap = BodyMocap(args.checkpoint_body_smplx, args.smpl_dir, device=device, use_smplx=True)
                hand_mocap = HandMocap(args.checkpoint_hand, args.smpl_dir, device=device)
                run_frank_mocap(args, hand_bbox_detector, body_mocap, hand_mocap, visualizer)
                
            self.status_var.set(f"处理完成! 输出保存到: {output_path}")
            
        except Exception as e:
            self.status_var.set(f"处理出错: {str(e)}")
            import traceback
            traceback.print_exc()

def main():
    # 设置主题
    root = tk.Tk()
    app = UnifiedMocapApp(root)
    root.mainloop()

if __name__ == "__main__":
    main()